#include "framework.h"

#include <initguid.h>
#include <evr.h>

#include <d3d9.h>
#include <dvdmedia.h>
#include <dxva2api.h>

#include <mfapi.h>
#include <mferror.h>
#include <mftransform.h>

IDirect3DDeviceManager9* deviceManager;

IMFTransform* mfTransform;
int width;
int height;
unsigned long samplesize;
unsigned long subtype;

IMFSample* dstSample;

HRESULT initialize(IMFTransform* transform)
{
	HRESULT result = S_OK;

	IMFMediaType* type = NULL;
	IMFMediaBuffer* dstBuffer = NULL;
	IDirect3DSurface9* dstSurface = NULL;
	HANDLE hDevice = NULL;
	IDirectXVideoAccelerationService* service = NULL;
	IDirect3DDevice9* pDevice = NULL;

	AM_MEDIA_TYPE* format = NULL;
	VIDEOINFOHEADER2* header = NULL;

	// Release what we already had beforehand
	SAFE_RELEASE(mfTransform);
	SAFE_RELEASE(dstSample);

	ASSERT(transform->GetOutputCurrentType(0, &type));
	ASSERT(type->GetRepresentation(AM_MEDIA_TYPE_REPRESENTATION, (void**)&format));

	header = (VIDEOINFOHEADER2*)format->pbFormat;

	width = header->bmiHeader.biWidth;
	height = header->bmiHeader.biHeight;
	samplesize = format->lSampleSize;
	subtype = format->subtype.Data1;

	// Video memory sample
	ASSERT(deviceManager->OpenDeviceHandle(&hDevice));
	ASSERT(deviceManager->LockDevice(hDevice, &pDevice, FALSE));
	ASSERT(pDevice->CreateOffscreenPlainSurface(width, height, (D3DFORMAT)subtype, D3DPOOL_DEFAULT, &dstSurface, NULL));
	ASSERT(MFCreateDXSurfaceBuffer(IID_IDirect3DSurface9, dstSurface, FALSE, &dstBuffer));
	ASSERT(MFCreateVideoSampleFromSurface(NULL, &dstSample));
	ASSERT(dstSample->AddBuffer(dstBuffer));
	ASSERT(deviceManager->UnlockDevice(hDevice, FALSE));
	ASSERT(deviceManager->CloseDeviceHandle(hDevice));

	mfTransform = transform;
	mfTransform->AddRef();

end:
	if (format)
		type->FreeRepresentation(AM_MEDIA_TYPE_REPRESENTATION, format);

	SAFE_RELEASE(service);
	SAFE_RELEASE(dstBuffer);
	SAFE_RELEASE(dstSurface);
	SAFE_RELEASE(type);

	return result;
}

VTABLE_HOOK(HRESULT, IMFTransform, ProcessOutput, DWORD dwFlags, DWORD cOutputBufferCount,
	MFT_OUTPUT_DATA_BUFFER* pOutputSamples, DWORD* pdwStatus)
{
	if (pOutputSamples->pSample != NULL || deviceManager == NULL || cOutputBufferCount != 1)
		return originalProcessOutput(This, dwFlags, cOutputBufferCount, pOutputSamples, pdwStatus);

	HRESULT result = S_OK;

	if (mfTransform != This)
		ASSERT(initialize(This));

	pOutputSamples->pSample = dstSample;
	pOutputSamples->pSample->AddRef();

	ASSERT(originalProcessOutput(This, dwFlags, cOutputBufferCount, pOutputSamples, pdwStatus));

end:
	return result;
}

VTABLE_HOOK(HRESULT, IMFTransform, ProcessMessage, MFT_MESSAGE_TYPE eMessage, ULONG_PTR ulParam)
{
	HRESULT result = originalProcessMessage(This, eMessage, ulParam);
	if (eMessage == MFT_MESSAGE_SET_D3D_MANAGER && result == MF_E_UNSUPPORTED_D3D_TYPE)
	{
		result = originalProcessMessage(This, eMessage, NULL);
		if (SUCCEEDED(result))
			INSTALL_VTABLE_HOOK(This, ProcessOutput, 25);
	}
	return result;
}

HOOK(void*, IMFTransformInitializer, 0x140420B90, void* a1, void* a2, IMFTransform** transform)
{
	void* result = originalIMFTransformInitializer(a1, a2, transform);
	INSTALL_VTABLE_HOOK(*transform, ProcessMessage, 23);
	return result;
}

HOOK(HRESULT, DXVA2CreateDirect3DDeviceManager, PROC_ADDRESS("dxva2.dll", "DXVA2CreateDirect3DDeviceManager9"),
	UINT* pResetToken, IDirect3DDeviceManager9** ppDeviceManager)
{
	HRESULT result = originalDXVA2CreateDirect3DDeviceManager(pResetToken, ppDeviceManager);

	if (FAILED(result))
		return result;

	if (deviceManager)
		deviceManager->Release();

	deviceManager = *ppDeviceManager;
	deviceManager->AddRef();

	return result;
}

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
	if (ul_reason_for_call == DLL_PROCESS_ATTACH)
	{
		INSTALL_HOOK(DXVA2CreateDirect3DDeviceManager);
		INSTALL_HOOK(IMFTransformInitializer);
	}

	return TRUE;
}